import collections
import getpass
import os
import subprocess
import sys
import tempfile
import time

from batchcat.exceptions import JobSubmitFail
from batchcat import jobscripts

class ResourceManager(object):
    def __init__(self, name, default_queue, unthrottled_limit):
        self.name = name
        self.queue = default_queue
        self.unthrottled_limit = unthrottled_limit

        self._joblist = []

        self.pause = 5

    def submit(self):
        self._joblist.reverse()
        njobs_total = len(self._joblist) 

        while len(self._joblist) > 0:
            jobs_running = self.running_job_count()

            njobs_to_submit = self.unthrottled_limit - jobs_running
            if njobs_to_submit > len(self._joblist):
                njobs_to_submit = len(self._joblist)

            for i in xrange(njobs_to_submit):
                jobscr = self._joblist.pop()
                with tempfile.NamedTemporaryFile() as fp:
                    fp.write(str(jobscr))
                    fp.flush()
                    
                    self._submit(fp.name)

            time.sleep(self.pause)                
            sys.stdout.write("{0}/{1} jobs submitted\r".format(njobs_total - len(self._joblist),
                                                               njobs_total))
            sys.stdout.flush()


        print "{0}/{0} jobs submitted\r".format(njobs_total - len(self._joblist),
                                                njobs_total)
        return njobs_total

    def _submit(self, path):
        pass

    def running_job_count(self):
        raise NotImplementedError("Abstract base class not for end use")

    def add_job(self):
        pass
    
    def add_script(self, scrobj):
        self._joblist.append(scrobj)
    
    def write(self, path=None):
        if path is None:
            path = os.getcwd()
        elif os.access(path, os.F_OK | os.W_OK):
            if not os.path.isdir(path):
                raise IOError("Cannot write scripts: path `{0}' exists and is not a directory".format(path))
        else:
            os.mkdir(path)

        tot_count = 0
        
        nscripts = len(self._joblist)
        for script_obj in self._joblist:
            if tot_count % 100 == 0:
                sys.stdout.write("{0}/{1} scripts written\r".format(tot_count, nscripts))
                sys.stdout.flush()

            with open("{0}.{1}x.sh".format(os.path.join(path, script_obj.name), 1), 'w') as fp:
                fp.write(str(script_obj))

            tot_count += 1

        print "{0}/{0} scripts written".format(nscripts)

        return tot_count


class SlurmResourceManager(ResourceManager):
    def __init__(self, name, default_queue, unthrottled_limit=1000):
        super(SlurmResourceManager, self).__init__(name, default_queue, unthrottled_limit)

    def _submit(self, script_path):
        rm_pipe = subprocess.Popen('sbatch -p {0}'.format(self.queue).split() + [script_path],
                                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        stdout, stderr = rm_pipe.communicate()
                
        if rm_pipe.returncode != 0:
            msg = "Error submitting jobs through `{0}'".format(self.name)
            msg = "Resource manager output:"
            msg += "\n\tstdout: `{0}'".format(stdout)
            msg += "\n\tstderr: `{0}'".format(stderr)
            
            raise JobSubmitFail(msg)

    def running_job_count(self):
        username = getpass.getuser()

        cmdargs = 'squeue -a -u {0}'.format(username)

        squeue_pipe = subprocess.Popen(cmdargs.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        stdout, stderr = squeue_pipe.communicate()

        if squeue_pipe.returncode != 0:
            sys.stderr.write("Subprocess call failed -- `{0}':".format(cmdargs))
            sys.stderr.write(str(stderr) + '\n')
            sys.stderr.write("Aborting...\n")
            raise OSError

        lines = [line for line in stdout.split('\n') if username in line]

        return len(lines)
    
    def add_job(self, name, command, opts):
        self._joblist.append(jobscripts.SlurmScript(name, command, opts))




resource_managers = {'olympus': SlurmResourceManager('olympus', 'shared'),
                     }

resource_managers['default'] = resource_managers['olympus']


def set_default_rm(name):
    if name in resource_managers:
        resource_managers['default'] = resource_managers['sbatch']
    else:
        raise KeyError("Unrecognized resource manager name `{0}'".format(name))

 
